import os
from shutil import copyfile
import shutil
import math

def find_images(path):
    images_name = list()
    valid_extensions = ['.jpg', '.jpeg', '.png', '.bmp']
    for root, dirs, files in os.walk(path):
        for file_name in files:
            _, extension = os.path.splitext(file_name)
            if extension.lower() in valid_extensions:
                images_name.append(file_name)

    return images_name


import os
import re

def rename_pmi(folder_path):
    index = 0
    # 遍历文件夹中的文件
    for filename in os.listdir(folder_path):
        # 提取文件名和扩展名
        basename, extension = os.path.splitext(filename)
        # new_file_name = f"{index}{extension}"

        # new_file_path = os.path.join(folder_path, new_file_name)
        # index += 1
        # 使用正则表达式提取数字和括号中的内容
        pattern = r"\((.*?)\)"
        match = re.search(pattern, basename)
        if match:
            content = match.group(1)  # 提取括号中的内容
            content = re.sub(r"\D", "_", content)  # 将非数字字符替换为下划线

            # 构建新的文件名
            new_filename = basename.replace(match.group(0), content)
            new_filename = new_filename.replace(" ", "_")
            new_filename += extension

            # 构建新的文件路径
            new_filepath = os.path.join(folder_path, new_filename)

            # 重命名文件
            os.rename(os.path.join(folder_path, filename), new_filepath)

        # 重命名文件
        # os.rename(os.path.join(folder_path, filename), new_file_path)


def rename_aoi(folder_path):
    # 遍历文件夹中的文件
    for filename in os.listdir(folder_path):
        # 提取文件名和扩展名
        basename, extension = os.path.splitext(filename)
        # new_file_name = f"{index}{extension}"

        # new_file_path = os.path.join(folder_path, new_file_name)
        # index += 1
        # 使用正则表达式提取数字和括号中的内容
        pattern = r"[▒▒╕╛]"
        new_basename = re.sub(pattern, "", basename)
        
        new_filename = new_basename + extension

        # 构建新的文件路径
        new_filepath = os.path.join(folder_path, new_filename)

        # 重命名文件
        os.rename(os.path.join(folder_path, filename), new_filepath)


def gen_ini(
    images_dir: str, # image data path
    ini_file: str = "/data/lot/ColorImageGrabingInfo.ini", # ColorImageGrabingInfo.ini file
    pixel_size: int = 1, # change pixel_size
):
    """
    [97958.42860.c.-1927802918.jpeg]
    X=97960.2871311362
    Y=42856.3340620407
    FaultX=97958.2589232393
    FaultY=42857.5626754375
    PixelSizeX=0.46398800611496
    PixelSizeY=0.464109987020493
    ImageSizeX=1024
    ImageSizeY=768
    Col=95
    Row=34
    Type=0
    Index=280
    IsTableCoordinate=0
    IsColor=1
    IsGreyRevisit=0
    IsBsi=0
    IsClip=0
    IsClip2=0
    """
    # 图片个数
    images_name = find_images(images_dir)

    # 源 ini 行数
    with open(ini_file, "r") as fr:
        lines = fr.readlines()

    fake_ini = os.path.join(images_dir, "ColorImageGrabingInfo.ini")

    if len(images_name) < len(lines)//19:
        img_len = len(images_name)   
    else:              
        img_len = math.ceil(len(images_name) / (len(lines)//19)) * len(lines)//19

    with open(fake_ini, "w") as fw:
        for i in range(img_len*19):
            if lines[i].startswith("[") and lines[i].endswith("]\n"):
                new_line = "[" + images_name[i//19] + "]\n"
            elif "PixelSizeX" in lines[i]:
                new_line = "PixelSizeX=" + str(pixel_size) + "\n"
            elif "PixelSizeY" in lines[i]:
                new_line = "PixelSizeY=" + str(pixel_size) + "\n"
            else:
                new_line = lines[i]
            fw.writelines(new_line)


    # copyfile("/data/lot/ProductInfo.ini", os.path.join(images_dir, "ProductInfo.ini"))
    # copyfile("/data/lot/WaferMapRecipe.ini", os.path.join(images_dir, "WaferMapRecipe.ini"))


def copy_allfiles(src,dest):
    #src:原文件夹；dest:目标文件夹

    # 如果目标目录不存在，则创建它
    if not os.path.exists(dest):
        os.makedirs(dest)
    src_files = os.listdir(src)
    for file_name in src_files:
        full_file_name = os.path.join(src, file_name)
        if os.path.isfile(full_file_name):
            shutil.copy(full_file_name, dest)


def gen_aoi(images_dir):
    
    rename_aoi(images_dir)
    gen_ini(images_dir)


if __name__ == "__main__":
    pmi = False
    aoi = True

    # 存放图片目录的名字
    aoi_lotname = "12-E5"
    pmi_lotname = "FQR1FR800XTNT4DU"
    wafer_index = "01"
    dir_folder = "test_lot"
    # dir_folder = "all_lot"

    new_aoi_lotname = aoi_lotname.replace('-', '')
    pmi_lotname = pmi_lotname.replace('-', '')

    if pmi:
        pmi_images_dir = f"/data/lot/{pmi_lotname}"
        pmi_wafer_dir = f"/data/lot/{dir_folder}/{pmi_lotname}/{wafer_index}"
        txt_file = f"/data/lot/{dir_folder}/{pmi_lotname}/{pmi_lotname}-{wafer_index}-B1.txt"

        # rename_pmi(pmi_images_dir)

        copy_allfiles(pmi_images_dir, pmi_wafer_dir)
        shutil.copy("/data/lot/AP41934-01-B1.txt", txt_file)

        run(pmi_wafer_dir)

    if aoi:
        aoi_images_dir = f"/data/lot/{aoi_lotname}"
        aoi_wafer_dir = f"/data/lot/{dir_folder}/{new_aoi_lotname}/{wafer_index}"
        txt_file = f"/data/lot/{dir_folder}/{new_aoi_lotname}/{new_aoi_lotname}-{wafer_index}-B1.txt"

        # rename_aoi(aoi_images_dir)

        # 拷贝图片到wafer目录
        copy_allfiles(aoi_images_dir, aoi_wafer_dir)
        shutil.copy("/data/lot/AP41934-01-B1.txt", txt_file)

        run(aoi_wafer_dir)


